package belf.migrate.engine.job.structure.helper;

import belf.migrate.api.util.EngineConf;
import belf.migrate.core.util.StringTool;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import com.google.gson.Gson;
import com.google.gson.JsonObject;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 使用LLM大模型对SQL进行转换
 */
@Slf4j
public class SqlAgent {
    // 创建 OkHttpClient 对象
    private static OkHttpClient client = new OkHttpClient.Builder()
            .connectTimeout(30, TimeUnit.SECONDS) // 连接超时时间，这里设置为 10 秒
               .readTimeout(45, TimeUnit.SECONDS) // 读取超时时间，这里设置为 15 秒
               .writeTimeout(30, TimeUnit.SECONDS) // 写入超时时间，这里设置为 15 秒
               .build();

    private static String endpoint = EngineConf.getInstance().getString("llm.server.endpoint");

    /**
     * 各个SQL转化调用组件，自己拼凑合适的promt，要求输出格式为xml，包含一个“sql”节点和一个“reasons”节点
     *
     * @param model  使用哪个大模型
     * @param prompt 提示词
     * @return 转化后的sql和转换理由
     */
    public static SqlResponse post(String model, String prompt) {
        String template = """
                {
                    "model": "%s",
                    "messages": [
                        {
                            "role": "system",
                            "content": "你是一名经验丰富的DBA."
                        },
                        {
                            "role": "user",
                            "content": "%s"
                        }
                    ]
                }
                """;
        // 使用 Gson 创建 JSON 对象
        JsonObject rootObject = new JsonObject();
        rootObject.addProperty("model", model);

        JsonArray messages = new JsonArray();
        JsonObject systemNode = new JsonObject();
        systemNode.addProperty("role", "system");
        systemNode.addProperty("content", "你是一名经验丰富的DBA.");
        messages.add(systemNode);

        JsonObject useNode = new JsonObject();
        useNode.addProperty("role", "user");
        useNode.addProperty("content", prompt);
        messages.add(useNode);

        rootObject.add("messages", messages);
        // llm的temperature参数越低，输出越接近真实情况
        rootObject.addProperty("temperature", EngineConf.getInstance().getFloat("llm.temperature", 0.1f));
        // llm的top_p参数越低，输出越接近真实情况。一般temperature和top_p参数二选一即可
        rootObject.addProperty("top_p", EngineConf.getInstance().getFloat("llm.temperature", 0.3f));
        // 输出 JSON 字符串
        Gson gson = new Gson();
        String json = gson.toJson(rootObject);
        // 设置请求体为 JSON 格式
        RequestBody requestBody = RequestBody.create(json, MediaType.parse("application/json"));

        SqlResponse sqlResponse = null;
        try {
            // 创建 POST 请求
            Request request = new Request.Builder()
                    .url(endpoint)
                    .post(requestBody)
                    .build();
            // 发送请求并获取响应
            Response response = client.newCall(request).execute();
            // 打印响应内容
            if (response.isSuccessful()) {
                sqlResponse = parseChatgptResponse(response.body().string());
            }
            response.close();
        } catch (Exception e) {
            log.error("调用大模型转换SQL失败：{}", e.getMessage());
            e.printStackTrace();
        }

        return sqlResponse;
    }

    public static void close() {
        // 关闭客户端连接
        client.dispatcher().executorService().shutdown();
        client.connectionPool().evictAll();
    }

    private static SqlResponse parseChatgptResponse(String jsonString) {
        // 解析Chatgpt输出 JSON 字符串
        Gson gson = new Gson();
        JsonObject rootObject = gson.fromJson(jsonString, JsonObject.class);
        JsonArray choices = rootObject.getAsJsonArray("choices");
        if (null != choices && choices.size() > 0) {
            JsonObject jsonObject = choices.get(0).getAsJsonObject();
            JsonObject message = jsonObject.get("message").getAsJsonObject();
            String content = message.get("content").getAsString();
//            System.out.println("LLM response content: \n" + content);
            String sql = StringTool.substring(content, "```sql", "```").trim();
            sql = StringTool.trimRight(sql, "/");
            String reasonStr = StringTool.lastSubstring(content, "```").trim();
            List<String> reasons = Arrays.asList(reasonStr.split("\n"));
            return new SqlResponse(sql, reasons);
        }
        return null;
    }

    private static SqlResponse parseResponse(String xmlString) {
        String sql = "";
        List<String> reasons = new ArrayList<>();
        try {
            DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
            DocumentBuilder builder = factory.newDocumentBuilder();
            Document document = builder.parse(new org.xml.sax.InputSource(new java.io.StringReader(xmlString)));

            // 获取根元素
            Element rootElement = document.getDocumentElement();

            // 获取 sql 节点内容
            NodeList sqlNodes = rootElement.getElementsByTagName("sql");
            if (sqlNodes.getLength() > 0) {
                Node sqlNode = sqlNodes.item(0).getFirstChild();
                if (sqlNode != null) {
                    sql = sqlNode.getNodeValue();
                }
            }

            // 获取 reasons 节点下的所有 reason 节点内容
            NodeList reasonNodes = rootElement.getElementsByTagName("reasons").item(0).getChildNodes();
            for (int i = 0; i < reasonNodes.getLength(); i++) {
                Node reasonNode = reasonNodes.item(i);
                reasons.add(reasonNode.getTextContent());
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
        return new SqlResponse(sql, reasons);
    }

    public static List<String> getModels() {
        return null;
    }

}
